import numpy as np


class BaseEnv(object):
    """
    基础环境类 - 强化学习环境的抽象基类

    这个类定义了所有强化学习环境必须实现的基本接口。
    它遵循OpenAI Gym的API设计，提供了一个标准的环境交互框架。

    环境的基本概念：
    - 智能体（Agent）：在环境中执行动作的实体
    - 状态（State/Observation）：环境当前情况的描述
    - 动作（Action）：智能体可以执行的操作
    - 奖励（Reward）：环境对智能体动作的反馈
    - 完成（Done）：标志当前回合是否结束

    环境的使用流程：
    1. reset() -> 重置环境，获得初始状态
    2. step(action) -> 执行动作，获得(下一状态, 奖励, 完成标志, 信息)
    3. 如果完成标志为True，回到步骤1
    """

    def __init__(self, rank) -> None:
        """
        初始化基础环境

        Args:
            rank (int): 环境的排名或ID，用于区分多个并行环境实例
                         在多进程或多线程训练中，每个环境实例都有唯一的rank
        """
        # 观察空间：定义智能体可以观察到的信息格式和范围
        # 具体实现由子类定义，例如可以是图像、向量、字典等
        self.observation_space = None

        # 动作空间：定义智能体可以执行的动作类型和范围
        # 具体实现由子类定义，例如可以是离散动作、连续动作等
        self.action_space = None

        # 环境排名/ID，用于在并行环境中区分不同的环境实例
        self.rank = rank

    def step(self, act):
        """
        执行一步环境交互

        这是环境与智能体交互的核心方法。智能体提供一个动作，
        环境根据这个动作执行一步模拟，并返回结果。

        Args:
            act: 智能体执行的动作，具体格式由action_space定义

        Returns:
            tuple: (obs, reward, done, info)
                - obs (torch.Tensor or np.ndarray): 新的观察状态
                    形状为(n_agent, ...)，其中n_agent是智能体数量
                - reward (torch.Tensor or np.ndarray): 奖励信号
                    形状为(n_agent, 1) 或 (n_team, 1)
                    n_agent：每个智能体单独获得奖励
                    n_team：每个队伍共享奖励（团队奖励）
                - done (bool): 回合结束标志
                    True表示当前回合结束，需要调用reset()
                - info (dict): 额外的调试信息，可能包含：
                    - 环境状态详情
                    - 性能指标
                    - 调试数据
                    - 其他相关信息

        Raises:
            NotImplementedError: 子类必须实现此方法

        注意：
        - 如果只有一个队伍且使用统一奖励(TakeRewardAsUnity)，
          必须确保奖励的shape=[n_team=1, 1]
          例如：
          >> RewardForTheOnlyTeam = +1
          >> RewardForAllTeams = np.array([RewardForTheOnlyTeam, ])
          >> return (ob, RewardForAllTeams, done, info)  # 统一奖励格式
          >> return (ob, RewardForAllAgents, done, info)  # 个别奖励格式
        """
        raise NotImplementedError

    def reset(self):
        """
        重置环境到初始状态

        当一个回合结束（done=True）或需要开始新回合时调用此方法。
        环境将回到初始状态，并返回初始观察。

        Returns:
            tuple: (obs, info)
                - obs (torch.Tensor or np.ndarray): 初始观察状态
                    形状为(n_agent, ...)，其中n_agent是智能体数量
                - info (dict): 重置时的额外信息，可能包含：
                    - 初始环境配置
                    - 智能体初始状态
                    - 种子信息
                    - 其他相关信息

        Raises:
            NotImplementedError: 子类必须实现此方法
        """
        raise NotImplementedError


class RawObsArray(object):
    """
    原始观察数组类

    这个类用于动态管理原始观察数据的存储格式和大小。
    它支持在不知道最终数据大小的情况下逐步添加数据，
    然后在第一次调用get()时确定最终的数据大小。

    主要用途：
    - 处理变长的观察数据
    - 自动推断观察数据的大小
    - 缓存观察数据以提高性能

    使用场景：
    - 某些环境的观察数据长度可能会变化
    - 需要在运行时确定观察数据的格式
    - 批量处理观察数据
    """

    # 类变量：所有实例共享的原始观察大小字典
    # 键：观察数据的key，值：对应的数据大小
    raw_obs_size = {}

    def __init__(self, key='default'):
        """
        初始化原始观察数组

        Args:
            key (str): 观察数据的标识符，用于区分不同类型的观察数据
                       默认为'default'，可以自定义如'image', 'vector', 'audio'等
        """
        self.key = key  # 观察数据的标识符

        # 检查是否已经知道这种观察数据的大小
        if self.key not in self.raw_obs_size:
            # 未知大小：使用列表暂存数据，等待第一次get()调用时确定大小
            self.guards_group = []  # 临时存储数据的列表
            self.nosize = True      # 标记：大小未知
        else:
            # 已知大小：直接创建固定大小的numpy数组
            size = self.raw_obs_size[self.key]
            self.guards_group = np.zeros(shape=(size), dtype=np.float32)
            self.nosize = False     # 标记：大小已知
            self.p = 0              # 当前写入位置指针

    def append(self, buf):
        """
        向观察数组中添加数据

        根据是否已知数据大小，采用不同的存储策略。

        Args:
            buf (np.ndarray): 要添加的观察数据缓冲区
                              可以是单个观察或批量观察
        """
        if self.nosize:
            # 大小未知：将数据添加到列表中暂存
            self.guards_group.append(buf)
        else:
            # 大小已知：直接写入预分配的numpy数组
            L = len(buf)  # 获取数据长度
            # 将数据写入当前指针位置
            self.guards_group[self.p:self.p + L] = buf[:]
            # 更新指针位置
            self.p += L

    def get(self):
        """
        获取完整的观察数组

        如果是第一次调用且大小未知，则会：
        1. 将列表中的所有数据合并为一个numpy数组
        2. 记录数据大小到类变量中
        3. 返回合并后的数组

        Returns:
            np.ndarray: 完整的观察数据数组

        注意：
        - 第一次调用后会改变内部存储结构，从列表变为numpy数组
        - 后续调用会直接返回numpy数组，提高访问效率
        """
        if self.nosize:
            # 第一次调用且大小未知：合并列表中的数据
            self.guards_group = np.concatenate(self.guards_group)
            # 记录数据大小到类变量，供后续实例使用
            self.raw_obs_size[self.key] = len(self.guards_group)
            # 标记大小已知
            self.nosize = False
        return self.guards_group

    def get_raw_obs_size(self):
        """
        获取原始观察数据的大小

        Returns:
            int: 观察数据的长度

        Raises:
            AssertionError: 如果不知道观察数据大小

        注意：
        - 必须在调用过get()方法后才能调用此方法
        - 或者通过其他方式已经确定了数据大小
        """
        assert self.key in self.raw_obs_size, f"不知道观察数据'{self.key}'的大小"
        return self.raw_obs_size[self.key]


# 使用示例和说明：
"""
基础环境类的使用示例：

1. 创建自定义环境：
class MyEnvironment(BaseEnv):
    def __init__(self, rank):
        super().__init__(rank)
        self.observation_space = gym.spaces.Box(low=0, high=1, shape=(10,))
        self.action_space = gym.spaces.Discrete(4)
        self.state = np.random.random(10)
        self.steps = 0

    def step(self, act):
        # 执行动作
        self.state = np.random.random(10)
        self.steps += 1

        # 计算奖励
        reward = np.random.random()

        # 检查是否结束
        done = self.steps >= 100

        # 返回结果
        return self.state, reward, done, {"steps": self.steps}

    def reset(self):
        self.state = np.random.random(10)
        self.steps = 0
        return self.state, {"reset": True}

2. 使用环境：
env = MyEnvironment(rank=0)

# 重置环境
obs, info = env.reset()

# 交互循环
for step in range(200):
    action = env.action_space.sample()  # 随机动作
    obs, reward, done, info = env.step(action)

    if done:
        print(f"回合结束，总步数：{step + 1}")
        obs, info = env.reset()

原始观察数组的使用示例：

1. 处理未知大小的数据：
obs_array = RawObsArray(key='vector')

# 逐步添加数据
obs_array.append(np.array([1, 2, 3]))
obs_array.append(np.array([4, 5, 6]))
obs_array.append(np.array([7, 8, 9]))

# 第一次获取：会合并数据并确定大小
full_obs = obs_array.get()
print(full_obs.shape)  # (9,)
print(obs_array.get_raw_obs_size())  # 9

2. 使用已知大小的数据：
# 第二个实例会自动使用已知的大小
obs_array2 = RawObsArray(key='vector')
obs_array2.append(np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]))
full_obs2 = obs_array2.get()  # 直接返回numpy数组

设计模式说明：
- BaseEnv使用了模板方法模式，定义了环境的基本接口
- RawObsArray使用了延迟初始化模式，在第一次使用时确定数据格式
- 这两个类共同构成了一个灵活的强化学习环境框架
"""